//-----------------------------------------------------------------------------
//
// NAT.c
//
// Group:
//    Doug Kline,          dakline@ncsu.edu
//    Mauricio Del Carpio, mgdelcar@ncsu.edu
//    Josh Jodrey,         jdjodrey@ncsu.edu
//
// This class handles the Network Address Translations.
//
//-----------------------------------------------------------------------------

#include <linux/string.h>
#include <linux/slab.h>
#include <linux/netfilter.h>
#include <linux/netfilter_ipv4.h>
#include <net/tcp.h>
#include <linux/semaphore.h>
//#include <linux/smp_lock.h>

#include "NAT.h"

extern struct Config g_config;

extern int debug;
extern int info;

static struct nf_hook_ops g_netfilter_ops_in;
static struct nf_hook_ops g_netfilter_ops_out;

static int g_nat_installed = 0;


struct Mapper
{
   unsigned int   host_address;
   unsigned int   server_address;
   unsigned short host_port;
   unsigned short server_port;
   
   struct Mapper *next_p;
};

//DEFINE_MUTEX(g_mutex);
DECLARE_MUTEX(g_mutex);

struct Mapper *g_map_head_p = NULL;


static void outgoing_map(struct iphdr *iph, struct tcphdr *tcph);
static void incoming_map(struct iphdr *iph, struct tcphdr *tcph);

static void lockNAT(void)
{
   down( &g_mutex );
}

static void unlockNAT(void)
{
   up( &g_mutex );
}

void print_ip(unsigned int ip, char* out_str)
{
    unsigned char ch1, ch2, ch3, ch4;
    ch1 = ip >> 24;
    ch2 = (ip >> 16) & 0xFF;
    ch3 = (ip >> 8) & 0xFF;
    ch4 = ip & 0xFF;
    
    sprintf(out_str, "%u.%u.%u.%u", ch1, ch2, ch3, ch4);
}

//----------------------------------------------------------------------------
//
// out_hook()
//
// Intercept outgoing packet to change the source id addr / tcp port number
//
//----------------------------------------------------------------------------
unsigned int out_hook(unsigned int hooknum,
                      struct sk_buff *skb,          /* socket buffer pointer */
                      const struct net_device *in,
                      const struct net_device *out,
                      int (*okfn)(struct sk_buff*))

{
   unsigned int orig_src;
   unsigned int orig_src_port;
   unsigned short orig_ip_checksum;
   unsigned short orig_tcp_checksum;
   char ip_dst[16];
   char ip_src1[16];
   char ip_src2[16];

   // See /usr/src/linux-headers-2.6.31.20-generic/include/net/ip.h
   // See /home/doug/linux-2.6.31/include/linux/ip.h
   struct iphdr  *iph;

   // See /usr/src/linux-headers-2.6.31.20-generic/include/net/tcp.h
   struct tcphdr *tcph;

   printk(KERN_DEBUG "START: out_hook()\n");

   if (!skb)
   {
      printk(KERN_DEBUG "END: out_hook() skb\n");
      return NF_ACCEPT;
   }

   iph = ip_hdr(skb);
   printk(KERN_DEBUG "After iph\n");
   if (!iph)
   {
      printk(KERN_DEBUG "END: out_hook() iph\n");
      return NF_ACCEPT;
   }
   
   tcph = (struct tcphdr *)((u_int32_t *)iph + iph->ihl);
   printk(KERN_DEBUG "After tcp\n");
   if (!tcph)
   {
      printk(KERN_DEBUG "END: out_hook() tcp\n");
      return NF_ACCEPT;
   }

   // Only need to run NAT on TCP packets for this project
   if (iph->protocol != IPPROTO_TCP)
   {
      printk(KERN_DEBUG "END: out_hook() protocol\n");
      return NF_ACCEPT;
   }

   orig_src          = iph->saddr;
   orig_src_port     = tcph->source;
   orig_ip_checksum  = iph->check;
   orig_tcp_checksum = tcph->check;
   
   // Do the NAT mapping
   outgoing_map(iph, tcph);
   printk(KERN_DEBUG "outgoing mapping\n");
   
   /* Now fix up the checksum */
   csum_replace4(&(iph->check), orig_src, iph->saddr);
	inet_proto_csum_replace4(&(tcph->check), skb, orig_src, iph->saddr, 1);
	inet_proto_csum_replace2(&(tcph->check), skb, orig_src_port, tcph->source, 0);
	
	if (info)
	{
	  print_ip(ntohl(orig_src), ip_src1);
	  print_ip(ntohl(iph->saddr), ip_src2);
	  print_ip(ntohl(iph->daddr), ip_dst);
	  	  	  
      printk(KERN_INFO "\nOutgoing packet\n");
      printk(KERN_INFO "source: %16s : %5d to %16s : %5d\n",
                         ip_src1,   ntohs(orig_src_port),
                         ip_src2, ntohs(tcph->source));
      printk(KERN_INFO "dest:   %16s : %5d\n", ip_dst, ntohs(tcph->dest));
//    printk(KERN_INFO "ip  checksum: %4x to %4x\n", orig_ip_checksum, iph->check);
//    printk(KERN_INFO "tcp checksum: %4x to %4x\n\n", orig_tcp_checksum, tcph->check);
   }

   printk(KERN_DEBUG "END: out_hook()\n");

   return NF_ACCEPT;

} // end out_hook()


//----------------------------------------------------------------------------
//
// in_hook()
//
// Intercept incoming packet to change the destination ip addr / tcp port number
//
//----------------------------------------------------------------------------
unsigned int in_hook(unsigned int hooknum,
                     struct sk_buff *skb,          /* socket buffer pointer */
                     const struct net_device *in,
                     const struct net_device *out,
                     int (*okfn)(struct sk_buff*))

{
   unsigned int orig_dest;
   unsigned int orig_dest_port;
   unsigned short orig_ip_checksum;
   unsigned short orig_tcp_checksum;
   char ip_dst1[16];
   char ip_dst2[16];
   char ip_src[16];
   
   
   struct net_device *eth0;
  
   // See /usr/src/linux-headers-2.6.31-20-generic/include/net/ip.h
   // See /home/doug/linux-2.6.31/include/linux/ip.h
   struct iphdr  *iph;

   // See /home/doug/linux-2.6.31/include/linux/tcp.h
   struct tcphdr *tcph;

   printk(KERN_DEBUG "START: in_hook()\n");

   if (!skb)
   {
      printk(KERN_DEBUG "END: in_hook() skb\n");
      return NF_ACCEPT;
   }

   eth0 = dev_get_by_name(&init_net, "eth0");   //TODO: Replace this with /proc/...
   
   printk(KERN_DEBUG "eth0\n");
   
   if (eth0 == skb->dev)
   {
      printk(KERN_DEBUG "END: in_hook() eth0\n");
      return NF_ACCEPT;
   }

   iph = ip_hdr(skb);
   printk(KERN_DEBUG "iph\n");

   if (!iph)
   {
      printk(KERN_DEBUG "END: in_hook() iph\n");
      return NF_ACCEPT;
   }
   
   tcph = (struct tcphdr *)((u_int32_t *)iph + iph->ihl);
   printk(KERN_DEBUG "tcphdr\n");
   
   if (!tcph)
   {
       printk(KERN_DEBUG "END: in_hook() tcph\n");
      return NF_ACCEPT;
   }
   
   // Only need to run NAT on TCP packets for this project
   if (iph->protocol != IPPROTO_TCP)
   {
       printk(KERN_DEBUG "END: in_hook() protocol\n");
      return NF_ACCEPT;
   }

   orig_dest         = iph->daddr;
   orig_dest_port    = tcph->dest;
   orig_ip_checksum  = iph->check;
   orig_tcp_checksum = tcph->check;

   // Packets can either come in from our attached hosts or from the remote
   // servers.  If it is our proxy's port number, then  let it go to the proxy.
   //if (ntohs(orig_dest_port) == g_config.proxyPort)
   if (orig_dest == g_config.proxyAddr.s_addr)
   {
      printk(KERN_DEBUG "###### END: in_hook() matched proxy port number\n");
      return NF_ACCEPT;
   }


   // Do the NAT mapping
   incoming_map(iph, tcph);
   printk(KERN_DEBUG "incoming mapping\n");
 
   /* Now fix up the checksum */
   csum_replace4(&(iph->check), orig_dest, iph->daddr);
	inet_proto_csum_replace4(&(tcph->check), skb, orig_dest, iph->daddr, 1);
	inet_proto_csum_replace2(&(tcph->check), skb, orig_dest_port, tcph->dest, 0);

   if(info)
   {
	  print_ip(ntohl(iph->saddr), ip_src);
	  print_ip(ntohl(orig_dest), ip_dst1);
	  print_ip(ntohl(iph->daddr), ip_dst2);
   
      printk(KERN_INFO "\nIncoming packet\n");
      printk(KERN_INFO "source: %16s : %5d\n", ip_src, ntohs(tcph->source));
      printk(KERN_INFO "dest:   %16s : %5d to %16s : %5d\n",
                         ip_dst1,  ntohs(orig_dest_port),
                         ip_dst2, ntohs(tcph->dest));
//    printk(KERN_INFO "ip  checksum: %4x to %4x\n", orig_ip_checksum, iph->check);
//    printk(KERN_INFO "tcp checksum: %4x to %4x\n\n", orig_tcp_checksum, tcph->check);
   }

  printk(KERN_DEBUG "END: in_hook()\n");

  return NF_ACCEPT;

} // end in_hook()





//-----------------------------------------------------------------------------
// natInit()
//
//-----------------------------------------------------------------------------
int natInit(void)
{
   int err;
   
   printk(KERN_DEBUG "Installing NAT\n");
   
   // See /usr/src/linux-headers-2.6.31-20-generic/include/linux/netfilter.h
   
   // Install outgoing NAT
   g_netfilter_ops_out.hook     = out_hook;
   g_netfilter_ops_out.pf       = PF_INET;
   g_netfilter_ops_out.hooknum  = NF_INET_POST_ROUTING;
   g_netfilter_ops_out.priority = NF_IP_PRI_FIRST;
   err = nf_register_hook(&g_netfilter_ops_out);
   if (err)
   {
       printk(KERN_ERR "#1 nf_register_hook() returned an error\n");
   }

   // Install incoming NAT
   g_netfilter_ops_in.hook     = in_hook;
   g_netfilter_ops_in.pf       = PF_INET;
   g_netfilter_ops_in.hooknum  = NF_INET_PRE_ROUTING;
   g_netfilter_ops_in.priority = NF_IP_PRI_FIRST;
   err = nf_register_hook(&g_netfilter_ops_in);
   if (err)
   {
       printk(KERN_ERR "#2 nf_register_hook() returned an error\n");
   }

   g_nat_installed = 1;
   
   return(0);
   
} // end natInit()



//-----------------------------------------------------------------------------
// natExit()
//
//-----------------------------------------------------------------------------
int natExit(void)
{
   printk(KERN_DEBUG "Uninstalling NAT\n");
   
   nf_unregister_hook(&g_netfilter_ops_in);
   nf_unregister_hook(&g_netfilter_ops_out);
   
   g_nat_installed = 0;

   return 0;
   
} // end natExit()


//-----------------------------------------------------------------------------
// natInstalled()
//
//-----------------------------------------------------------------------------
int natInstalled(void)
{
   return g_nat_installed;
}
   
   
//-----------------------------------------------------------------------------
// incoming_map()
//
// Maps the incoming destination address & port
//
// This function will only be called if the destination port is 80
//
//-----------------------------------------------------------------------------
void incoming_map(struct iphdr *iph, struct tcphdr *tcph)
{
   struct Mapper *map_p = g_map_head_p;
   struct Mapper *new_p;

   // Lock down the global data
   lockNAT();

   // Look up the port mapping if already exists
   
   while (map_p)
   {
      if ((iph->saddr == map_p->host_address) && (tcph->source == map_p->host_port))
      {
         //Match, update the server info in case it has changed
         map_p->server_address  = iph->daddr;
         map_p->server_port     = tcph->dest;

         goto incoming_map_exit;
      }
   
      if (map_p->next_p == NULL)
      {
         break;
      }
      
      map_p = map_p->next_p;
   }
   
   // If we came this far then we need to create one
   
   // TODO ??? Is this where we are locking up??? Try unlocking first???
   unlockNAT();
   new_p = kmalloc(sizeof(struct Mapper), GFP_KERNEL);
   
   if (new_p == NULL)
   {
//      unlockNAT();
      printk(KERN_ERR "ERROR: failed kmalloc call\n");
      return;
   }
   
   lockNAT(); // TODO ????
   
   if (g_map_head_p == NULL)
   {
      g_map_head_p = new_p; // First one
      map_p = new_p;
   }
   else
   {
      map_p->next_p = map_p; // Append to end
   }
   
   map_p->host_address    = iph->saddr;
   map_p->host_port       = tcph->source;
   map_p->server_address  = iph->daddr;
   map_p->server_port     = tcph->dest;
   map_p->next_p          = NULL;
   
   
incoming_map_exit:

   // Send the packet to our proxy server
   iph->daddr   = g_config.proxyAddr.s_addr;
   tcph->dest   = htons(g_config.proxyPort);

   unlockNAT();
   return;

} // end incoming_map()




//-----------------------------------------------------------------------------
// outgoing_map()
//
// Maps the outgoing source address/port
//
//-----------------------------------------------------------------------------
void outgoing_map(struct iphdr *iph, struct tcphdr *tcph)
{
   struct Mapper *map_p = g_map_head_p;
   
   // Lock down the global data
   lockNAT();

   // If the destination address/port matches an entry in our mapper table
   // then we need to fix up the source address/port
   
   while (map_p)
   {
      if ((iph->daddr == map_p->host_address) && (tcph->dest == map_p->host_port))
      {
         //Match
         iph->saddr = map_p->server_address;
         tcph->source = map_p->server_port;
         goto outgoing_map_exit;
      }
   
      if (map_p->next_p == NULL)
      {
         break;
      }
      
      map_p = map_p->next_p;
   }

outgoing_map_exit:
   unlockNAT();
   return;

} // end outgoing_map()



//-----------------------------------------------------------------------------
// lookupServerAddr()
//
// Looks up the server address that was previously stored in NAT.
//
// Returns: 1 if found, 0 if not found
//-----------------------------------------------------------------------------
int lookupServerAddr(unsigned int host_address, unsigned short host_port, 
                     unsigned int *server_address, unsigned short *server_port)
{
   struct Mapper *map_p = g_map_head_p;
   int retVal = 0;
   
   // Lock down the global data
   lockNAT();

   // If the destination address/port matches an entry in our mapper table
   // then we need to fix up the source address/port
   
   while (map_p)
   {
      if ((host_address == map_p->host_address) && (host_port == map_p->host_port))
      {
         //Match
         *server_address = map_p->server_address;
         *server_port    = map_p->server_port;
         retVal = 1;
         goto outgoing_map_exit;
      }
   
      if (map_p->next_p == NULL)
      {
         break;
      }
      
      map_p = map_p->next_p;
   }

outgoing_map_exit:
   unlockNAT();
   return retVal;



} // lookupServerAddr





